#include <opencv2/imgproc.hpp> // 包含图像处理功能
#include <opencv2/opencv.hpp>   // 包含 OpenCV 的核心功能
#include <opencv2/dnn.hpp>      // 包含 DNN 模块功能

using namespace cv; // 使用 OpenCV 命名空间
using namespace std;

// 函数：letterbox
// 目的：按照指定尺寸调整图像，并保持原始长宽比，用于目标检测输入
Mat letterbox(Mat &img, Size new_shape, Scalar color, bool _auto, bool scaleFill, bool scaleup, int stride)
{
    float width = img.cols; // 获取输入图像的宽度
    float height = img.rows; // 获取输入图像的高度
    
    // 计算缩放比例
    float r = min(new_shape.width / width, new_shape.height / height);
    
    // 如果不允许放大图像，限制缩放比例
    if (!scaleup)
        r = min(r, 1.0f);
    
    // 计算新的未填充的宽度和高度
    int new_unpadW = int(round(width * r));
    int new_unpadH = int(round(height * r));

    // 计算需填充的宽度和高度
    int dw = new_shape.width - new_unpadW;
    int dh = new_shape.height - new_unpadH;

    // 如果自动填充，确保填充的宽度和高度是 stride 的倍数
    if (_auto)
    {
        dw %= stride;
        dh %= stride;
    }

    // 将填充宽度和高度分为两半
    dw /= 2, dh /= 2;

    Mat dst; // 创建目标图像
    // 调整图像大小
    resize(img, dst, Size(new_unpadW, new_unpadH), 0, 0, INTER_LINEAR);
    
    // 计算每个方向的边界填充
    int top = int(round(dh - 0.1));
    int bottom = int(round(dh + 0.1));
    int left = int(round(dw - 0.1));
    int right = int(round(dw + 0.1));
    
    // 在图像周围添加边框，使用指定的颜色进行填充
    copyMakeBorder(dst, dst, top, bottom, left, right, BORDER_CONSTANT, color);
    
    return dst; // 返回处理后的图像
}

int main(int argc, char *argv[])
{
    // 输出当前使用的 OpenCV 版本
    std::cout << "OpenCV version: " << CV_VERSION << std::endl;

    // 指定 ONNX 模型的路径
    string model_path = "/home/lixinlong/Project/yolo-learn/yolov5s.onnx";
    
    // 读取 ONNX 模型
    dnn::Net model = dnn::readNetFromONNX(model_path);
    
    // 读取输入图像
    Mat img0 = imread("/home/lixinlong/Project/yolo-learn/Pictures/zidane.jpg");
    size_t row0 = img0.rows;
    size_t col0 = img0.cols;
    // 调整图像大小为640*640以适应模型输入
    Mat img;
    resize(img0, img, Size(640, 640));
    //Mat img = letterbox(img0, Size(640, 640), Scalar(114, 114, 114), true, false, true, 32);
    
    Mat blob;    
    // 将图像转换为 blob 格式
    dnn::blobFromImage(img, blob, 1 / 255.0f, Size(img.cols, img.rows), Scalar(0, 0, 0), true, false);
    
    // 设置 DNN 模型的后端和目标
    model.setPreferableBackend(dnn::DNN_BACKEND_OPENCV);
    model.setPreferableTarget(dnn::DNN_TARGET_CPU);
    
    // 设置输入 blob
    model.setInput(blob);
    
    // 获取输出层的名称
    vector<string> outLayerNames = model.getUnconnectedOutLayersNames();
    
    vector<Mat> result; // 用于存储模型输出结果
    
    // 执行前向推理
    model.forward(result, outLayerNames);

    // 获取输出数据中的第一层数据
    Mat out = Mat(result[0].size[1], result[0].size[2], CV_32F, result[0].ptr<float>());
    
    vector<Rect> boxes; // 用于存储检测到的框
    vector<int> indices; // 存储有效框的索引
    vector<float> scores; // 存储每个框的得分
    
    // 遍历输出数据，获取边框信息
    for (int r = 0; r < out.size[0]; r++)
    {
        float cx = out.at<float>(r, 0); // 计算中心 x 坐标
        float cy = out.at<float>(r, 1); // 计算中心 y 坐标
        float w = out.at<float>(r, 2);  // 获取宽度
        float h = out.at<float>(r, 3);  // 获取高度
        float sc = out.at<float>(r, 4);  // 获取置信度
        
        // 获取类别置信度
        Mat confs = out.row(r).colRange(5, 85);
        confs *= sc; // 更新置信度
        
        double minV = 0, maxV = 0; // 用于存储最小和最大值
        double *minI = &minV; // 最小值的指针
        double *maxI = &maxV; // 最大值的指针
        
        // 找到置信度的最大值及其索引
        minMaxIdx(confs, minI, maxI);
        scores.push_back(maxV); // 保存最大值得分
        boxes.push_back(Rect(cx - w / 2, cy - h / 2, w, h)); // 保存边框
        indices.push_back(r); // 保存索引
    }
    
    // 非极大值抑制，去除冗余框
    dnn::NMSBoxes(boxes, scores, 0.25f, 0.45f, indices);
    
    // 输出结果
    for (auto &ind : indices)
    {
        cout << indices[ind] << ":" << scores[ind] << endl; // 打印有效框的索引和得分
    }
}
